# -*- coding:utf-8 -*-
# editor: zzh
# date: 2022/10/13

from model import BertNer
from utills import tokenize, ids2tokens
import torch
import pickle
import json

class NerInfer(object):
    def __init__(self, model_check,dict, device = torch.device('cpu')):
        self.model = BertNer(num_class=len(dict) * 2)
        self.model.load_state_dict(torch.load(model_check, map_location='cpu'))
        self.model.to(device)
        self.dict = dict
        self.deivce = device
        self.reverse_dict = {self.dict[k]:k for k in self.dict}

    def ner(self,sentence):
        tokenized_ids, segment_ids, tokenized_tokens = tokenize(sentence)
        tokenized_ids = tokenized_ids.long().to(self.device)
        segment_ids = segment_ids.long().to(self.device)
        inputs = (tokenized_ids, segment_ids)
        ner_output, crf_output = self.model(inputs)
        tokens = ids2tokens(tokenized_ids[0])
        predcit_values = self.get_all_slots(crf_output[0], tokens)

    def get_all_slots(self,crf_output, tokens):
        res = []
        tmp = ""
        type = -1
        for i in range(len(crf_output)):
            if crf_output[i] == 0:
                if tmp != "" and type != -1:
                    res.append((self.reverse_dict(type), tmp))
                    type = -1
                    tmp = ""
            else:
                tmp += tokens[i]
                type = int((crf_output[i] + 1) / 2)

        if tmp != "" and type != -1:
            res.append((self.reverse_dict(type), tmp))
            type = -1
            tmp = ""

        return res



if __name__ == '__main__':

    model_check = ""
    dict = pickle.load(open(r'data/ents3ids.pkl', 'rb'))
    nerinfer = NerInfer(model_check, dict, device=torch.device('cuda:0'))


    test_data = json.load(open('data/1500.json','r',encoding='utf8'))

    results = []

    for dialogue in test_data:
        dia = []
        sents = dialogue['dialogue']
        for sent in sents:
            speaker = sent['speaker']
            text = sent['text']
            props = nerinfer.ner(text)
            dia.append({'speaker':speaker, 'text':text, 'prop':props})

        results.append({'id':sents['dialogue_idx'], 'dialgue':dia})

        if(len(results) == 100):
            break


    json.dump(results, open('ner_result.json','w',encoding='utf8'))



